import io
from typing import Annotated

import requests
from PIL import Image
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

import llms


@tool
def baiduSearchTool(text: str) -> str:
    """从百度搜索返回结果"""
    # str = requests.get("https://www.baidu.com/s?wd="+ text).text
    print("搜索关键字：", text)

    return "王创是思维拐点的老板"


class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]


# 构建一个图
graph_builder = StateGraph(State)

# tool = TavilySearchResults(max_results=2)
tool = baiduSearchTool
tools = [tool]
# tool.invoke("What's a 'node' in LangGraph?")

llm = llms.getLLM()
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}


# The first argument is the unique node name
# The second argument is the function or object that will be called whenever
# the node is used.
# 将节点添加到图
graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=[tool])
graph_builder.add_node("tools", tool_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.set_entry_point("chatbot")
graph = graph_builder.compile()

# 绘制图形成图片
im = Image.open(io.BytesIO(graph.get_graph().draw_mermaid_png()))
im.save("graph.png")
# im.show()

while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break
    for event in graph.stream({"messages": ("user", user_input)}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)
